# This class can be used to collect data from different sources.
# it can be the wireless antenna or a wired Ethernet interface
# for now we assume its only one


import socket
import dpkt
import sys
import pcap
import constants
import util
import hp

import re # Need this regular expression to search through 
import time
import datetime
import subprocess
import time

# Provides an interface to alleviate others from reading data
class CollectData:
    def __init__(self,iface,ifs_macs,wireless_ifs,freq,filename=None):
		self.iface = iface
		self.filename = filename # Name of the file where to read pcap data
		self.pcap_file = ""
		self.ifs_macs = ifs_macs#dictionary of ethernet interfaces and MACs
		self.wireless_ifs = wireless_ifs
		self.start_time = datetime.datetime.now()
		self.freq = freq
		self.hp = hp.HomeProfile()# TODO: Make this information addable/changeable, based on user action # Not just one time addition.
		self.hp.readDevInfo()

    # --
    # Opens and gives a handle of pcap file.
    # --
    def loadpcap(self):
        f = open(self.filename)
        self.pcap_file = dpkt.pcap.Reader(f)
    
    def getDevices(self):
        for ts, buf in self.pcap_file:
            eth = dpkt.ethernet.Ethernet(buf)
            src_mac = (eth.dst).encode("hex")
            dst_mac = (eth.src).encode("hex")
            #print src_mac
            #print dst_mac
            self.hp.add_mac(src_mac)
            self.hp.add_mac(dst_mac)
        print self.hp.get_macs()

    # --
    # This function is for development and debugging from pcap files.
    # --
    def printpcap(self):
        for ts, buf in self.pcap_file:
            print ts,len(buf)
            eth = dpkt.ethernet.Ethernet(buf)
            dst_mac = (eth.dst).encode("hex")
            src_mac = (eth.src).encode("hex")
            print "Source MAC:", util.add_colons_to_mac(src_mac)
            print "Dst MAC:",util.add_colons_to_mac(dst_mac)
            #print `dpkt.ethernet.Ethernet(pkt)`
            ip = eth.data
            print ip.p
            if(ip.p == dpkt.ip.IP_PROTO_TCP):
                tcp = ip.data
                print "source port:", tcp.sport
                print "dst port:", tcp.dport
            elif(ip.p == dpkt.ip.IP_PROTO_UDP):
                udp = ip.data
                print "source port:", udp.sport
                print "dst port:", udp.dport
                
            print "Source IP:", ip.src.encode("hex")
            print "Dst IP:",ip.dst.encode("hex")

    # --
    # Its only called once.
    # -- This function should get the list of asssociated clients with the AP
    # should get this information from the /proc/net/arp?
    # --
    def get_associated_clients(self):
        count = 0
        mac_regex = '([a-fA-F0-9]{2}[:]?){6}' 
        assoc_mac = ''
        try:
            lines = open(constants.PROC_NET_ARP).readlines()
        except IOError:
            print "Cannot open file", constants.PROC_NET_ARP
        for line in lines:
            if (count != 0):# we are the first line therefore skip it.
                assoc_list = line.split(" ")
                for item in assoc_list:
                    a = re.compile(mac_regex).search(item)
                    if a:
                        assoc_mac = item[a.start(): a.end()]
                        print assoc_mac
                pass
            count = count +1 
            self.hp.addAssocDevice(assoc_mac)
        pass

    # --
    # Set the MAC Addresses of the AP Intetrfaces
    # --
    def setAPMACs(self):
        print "Inside setAPMACs"
        for key,val in self.ifs_macs.iteritems():
            self.hp.addAPMAC(val)
            self.hp.add_mac(val)
        pass
        print "Total AP MACs seen: ",self.hp.get_ap_macs()
        print "Total MACs Seen: ", self.hp.get_macs()
    # --
    # Print IP Header Info
    # --
    def ip_hdr(self,ip):
        print "Source IP: ",ip.src.encode("hex")
        print "Destination IP: ",ip.dst.encode("hex")
        print "Proto Type: " , ip.p
        pass

    # --
    # This function is used to sniff from wire
    # TODO: Add Code for different formats of streams if required.
    # --
    def sniff(self):
        # For now sniffing on the Ethernet interface.
        # sniffing on "any" causes the packets to be received in cooked form
        # which looses the Ethernet frame information but gives rest of the information
        # For further details: http://wiki.wireshark.org/SLL or man page of packet 7
        pc = pcap.pcap(self.iface)
        #pc = pcap.pcap("eth0")
        print 'Listening on %s: With filter %s' % (pc.name, pc.filter)
        try:
            decode = {  pcap.DLT_LOOP:dpkt.loopback.Loopback,
                        pcap.DLT_NULL:dpkt.loopback.Loopback,
                        pcap.DLT_IEEE802:dpkt.ethernet.Ethernet,
                        pcap.DLT_EN10MB:dpkt.ethernet.Ethernet,
                        pcap.DLT_LINUX_SLL:dpkt.sll.SLL}[pc.datalink()]
        except KeyError:
            print pc.datalink()
            print "Please check if you are handling proper packet type"
        pass
        try:
            self.decode(pc)
            pass
        except KeyboardInterrupt:
            nrecv, ndrop, nifdrop = pc.stats()
            print '\n%d packets received by filter' % nrecv
            print '%d packets dropped by kernel' % ndrop
            
    def time_passed(self):
        #print "Inside Time Passed"
        cur_time = datetime.datetime.now()
        delta = cur_time - self.start_time
        #print self.freq
        if(delta.seconds > self.freq):
			return True	
        else:
			return False
	# --
	# This function is called to scan the devices
    # This function is here for reference as it will not work on AP. 
	# @args: 
	#		iface: Should be the interface on which we want to do scanning.
	# @retruns: 
	# 		retuurns output of the command
	# TODO:Add error handling for command existence and checking for hangups.
	# 
	# --
    def scan_devices_client(self):
        bssid = ""
        ssid = ""
        freq = 0
        sig_strength = 0
        channel = 0

        mac_regex = '([a-fA-F0-9]{2}[:]?){6}' 
		# Does scanning on interface self.iface
        scan_command = "iw dev wlan0 scan"
        f = subprocess.Popen(scan_command, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
        (f_stdin,f_stdout,f_stderr) = (f.stdin, f.stdout, f.stderr)
        if(constants.DEBUG_ENABLED):
            for i in f_stdout.readlines():
                print i
    	for i in f_stderr.readlines():
    	    if(i):
    	        if(constants.DEBUG_ENABLED):
    	            print i
    	            print 'ERROR:There was a commandline related error while running command', scan_command
    	        return False
    	    pass
    	pass
        for line in f_stdout.readlines():
            a = re.compile(mac_regex).search(line)
            if a:
                bssid = line[a.start(): a.end()]
                if(bssid.find(":") != -1):
                    print "MAC Found: ",bssid
                pass
            pass
            b = line.find("freq: ")
            if (b != -1):
                freq = int(line[6 : 11].strip())
                #print freq
            c = line.find("signal:")
            if (c != -1):
                sig_strength = float(line[len("signal:") + 1 : -4].strip())
                #print sig_strength
            d = line.find("SSID:")
            if (d != -1):
                ssid = line[len("SSID:")+1 : -1].strip()
                #print ssid
            e = line.find("DS Parameter set: channel")
            if (e != -1):
                channel = line[len("DS Parameter set: channel")+1 : -1].strip()
                #print channel
    
    # --
    # Function to remove devices from the self.hp.assoc_macs
    # --
    def remove_devices(self,refreshed_devices):
        #print "Inside remove_devices"
        assoc_macs = self.hp.get_assoc_macs() # Returns a list
        # Remove any device that left the AP
        for item in assoc_macs:
            if item not in refreshed_devices:
                self.hp.remAssocDevice(item)
                self.hp.logDevice(item,"left")
            pass
        pass
        #print "Associated Devices: " , self.hp.get_assoc_macs()
        
        
	# --
	# This function is called to get the station dump
	# @retruns: 
	# 		retuurns output of the command
	# TODO:Add error handling for command existence and checking for hangups.
	# 
	# --
    def scan_devices(self):
        #This is the list of macs that we found as a reuslt of running the command below
        refreshed_devices = [] 
        bssid = ""
        sig_avg = 0
        sig_strength = 0
        assoc_macs = self.hp.get_assoc_macs() # Returns a list

        mac_regex = '([a-fA-F0-9]{2}[:]?){6}' 
        scan_command = "iw dev " + self.iface +" station dump"
        f = subprocess.Popen(scan_command, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
        (f_stdin,f_stdout,f_stderr) = (f.stdin, f.stdout, f.stderr)
        if(constants.DEBUG_ENABLED):
            for i in f_stdout.readlines():
                print i
    	for i in f_stderr.readlines():
    	    if(i):
    	        if(constants.DEBUG_ENABLED):
    	            print i
    	            print 'ERROR:There was a commandline related error while running command', scan_command
    	        return False
    	    pass
    	pass
        for line in f_stdout.readlines():
            a = re.compile(mac_regex).search(line)
            if a:
                bssid = line[a.start(): a.end()]
                if(bssid.find(":") != -1):
                    refreshed_devices.append(bssid)# added the bssid
                    if bssid not in assoc_macs: # Adding the new found mac to associted devices to the AP
                        #print "New MAC Found: ",bssid
            	        self.hp.addAssocDevice(bssid)
                        self.hp.logDevice(bssid,"joined")
                pass
            pass
            b = line.find("signal avg: ")
            if (b != -1):
                sig_avg = int(line[6 : 11].strip())
                print "Signal Strength Avg: ",sig_avg
            c = line.find("signal:")
            if (c != -1):
                sig_strength = float(line[len("signal:") + 1 : -4].strip())
                #print "Signal Strength:", sig_strength
            pass
        pass
        self.remove_devices(refreshed_devices)





    # --
    # This function is used to decode the packets received from wire
    #   pc: pcap stream of packets
    # --
    def decode(self,pc):
        for ts, buf in pc:
            #print ts,len(buf)
            # Packet Debugging
            #print `decode(buf)`
            #eth = `dpkt.sll.SLL(buf)`
            eth = dpkt.ethernet.Ethernet(buf)
            #print `eth` # debug
            dst_mac = util.add_colons_to_mac((eth.dst).encode("hex"))
            src_mac = util.add_colons_to_mac((eth.src).encode("hex"))
            #print "Source MAC:", util.add_colons_to_mac(src_mac)
            #print "Dst MAC:",util.add_colons_to_mac(dst_mac)
            if(eth.type== dpkt.ethernet.ETH_TYPE_IP):
                ip = eth.data
                #self.ip_hdr(ip)
                if(ip.p == dpkt.ip.IP_PROTO_TCP):
                    tcp =ip.data
                    #print "source port:", tcp.sport
                    #print "dst port:", tcp.dport
                elif(ip.p == dpkt.ip.IP_PROTO_UDP):
                    udp =ip.data
                    #print "source port:", udp.sport
                    #print "dst port:", udp.dport
            if(eth.type== dpkt.ethernet.ETH_TYPE_ARP):
            	self.hp.add_mac(src_mac)
            	self.hp.addDevice(src_mac)
            if(self.time_passed()):
                self.start_time = datetime.datetime.now()
                self.scan_devices()#pass #perform scan command here	
                #print "Associated Devices: " , self.hp.get_assoc_macs()

                
